import cv2
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from yolo_msgs.msg import Mask
from cv_bridge import CvBridge
from ultralytics import YOLO
import numpy as np
from scipy.stats import mode

# 全局变量
ESC_KEY = 27
radius = 20  # 圆的半径
circle_color = (0, 0, 255)  # 圆的颜色（蓝色）
text_color = (0, 255, 0)  # 文本的颜色（绿色）
thickness = -1  # 设置为-1，表示要填充这个圆
font = cv2.FONT_HERSHEY_SIMPLEX  # OpenCV字体

# 加载YOLOv8模型
model = YOLO("yolov8n.pt")
from rclpy.qos import QoSProfile, QoSDurabilityPolicy

qos_profile = QoSProfile(
    depth=10,
    durability=QoSDurabilityPolicy.VOLATILE,
)
# 得到检测框的四个坐标
def xywh2abcd(x, y, w, h):
    output = list()
    x_min = int(max(x - w / 2, 0))
    x_max = int(min(x + w / 2, 640))  # 假设图像宽度为640
    y_min = int(max(y - h / 2, 0))
    y_max = int(min(y + h / 2, 480))  # 假设图像高度为480
    output.extend([x_min, y_min, x_max, y_max])
    return output

def cul_distance(depth_image, abcd):
    number_list = np.array([])  # 初始化为一个空的 NumPy 数组
    for i in range(abcd[1], abcd[3]):
        for j in range(abcd[0], abcd[2]):
            if depth_image[i, j] != 0:
                number_list = np.append(number_list, depth_image[i, j])

    valid_depths = number_list[number_list >= 150]

    if valid_depths.size == 0:
        return 0

    truncated_depths = np.floor(valid_depths / 100).astype(int)
    person_depth_mode, count = mode(truncated_depths)

    if person_depth_mode.size == 0:
        return 0

    return person_depth_mode[0] / 10.0

class ImageDepthProcessor(Node):
    def __init__(self):
        super().__init__('image_depth_processor')
        self.bridge = CvBridge()
        self.depth_image = None
        self.color_image = None

        # 创建订阅者来获取深度图像
        self.depth_subscription = self.create_subscription(
            Image,
            '/camera/depth/image_raw',
            self.depth_callback,
            10
        )

        # 创建订阅者来获取RGB图像
        self.color_subscription = self.create_subscription(
            Image,
            '/camera/color/image_raw',
            self.color_callback,
            10
        )

        # 创建发布者来发布目标位置和深度信息
        self.publisher = self.create_publisher(Mask, '/target_depth_info1', qos_profile)

        # 创建发布者来发布处理后的RGB图像
        self.annotated_image_publisher = self.create_publisher(Image, '/annotated_image', 10)

    def depth_callback(self, msg):
        # 将ROS深度图像消息转换为OpenCV格式
        self.depth_image = self.bridge.imgmsg_to_cv2(msg, "32FC1")

    def color_callback(self, msg):
        # 将ROS RGB图像消息转换为OpenCV格式
        self.color_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
        self.process_image()

    def process_image(self):
        if self.color_image is None or self.depth_image is None:
            return

        # 物体识别
        results = model.track(self.color_image, save=False, conf=0.4)
        boxes = results[0].boxes.xywh.cpu().numpy()  # 获取检测到的目标的边界框
        target_ids = results[0].boxes.cls.cpu().numpy()  # 获取目标类别
        scores = results[0].boxes.conf.cpu().numpy()  # 获取目标得分

        # 注释帧
        annotated_frame = results[0].plot()

        for box, target_id, score in zip(boxes, target_ids, scores):
            x, y, w, h = box
            center_x = int(x)
            center_y = int(y)

            # 在注释帧上画圆
            cv2.circle(annotated_frame, (center_x, center_y), radius, circle_color, thickness)

            # 获取检测框的四个坐标
            abcd = xywh2abcd(x, y, w, h)

            # 获取中心点的深度值
            if (0 <= center_x < self.depth_image.shape[1]) and (0 <= center_y < self.depth_image.shape[0]):
                depth_value = self.depth_image[center_y, center_x]
            else:
                depth_value = 888

            # 计算目标区域的深度众数
            target_depth = cul_distance(self.depth_image, abcd)

            # 创建要发布的消息
            depth_info = Mask()
            depth_info.target_id = int(target_id)
            depth_info.score = float(score)
            depth_info.x_min = int(abcd[0])
            depth_info.y_min = int(abcd[1])
            depth_info.x_max = int(abcd[2])
            depth_info.y_max = int(abcd[3])
            depth_info.center_x = int(center_x)
            depth_info.center_y = int(center_y)
            depth_info.target_depth = float(target_depth)
            print(depth_info)
            # 发布深度值和坐标
            self.publisher.publish(depth_info)
            self.get_logger().info(
                f'Published target depth info: id={target_id}, score={score}, bbox=({abcd[0]}, {abcd[1]}, {abcd[2]}, {abcd[3]}), center=({center_x}, {center_y}), depth={target_depth}')

            # 在图像上显示深度值
            depth_text = f"z: {target_depth}m"
            cv2.putText(annotated_frame, depth_text, (center_x, center_y - radius - 10), font, 0.5, text_color, 1,
                        cv2.LINE_AA)

        # 发布注释后的图像
        annotated_image_msg = self.bridge.cv2_to_imgmsg(annotated_frame, "bgr8")
        self.annotated_image_publisher.publish(annotated_image_msg)

def main(args=None):
    rclpy.init(args=args)
    image_depth_processor = ImageDepthProcessor()

    try:
        rclpy.spin(image_depth_processor)
    except KeyboardInterrupt:
        pass

    # 销毁节点
    image_depth_processor.destroy_node()
    rclpy.shutdown()

if __name__ == "__main__":
    main()
